#include <string.h>
#include <math.h>
#include "esp_log.h"
#include "color_predictor.h"

static const char *TAG = "ColorPredictor";

static float relu(float x) {
    return (x > 0) ? x : 0;
}

static void softmax(float* input, int size) {
    float max = input[0];
    for (int i = 1; i < size; i++) {
        if (input[i] > max) {
            max = input[i];
        }
    }

    float sum = 0;
    for (int i = 0; i < size; i++) {
        input[i] = expf(input[i] - max);
        sum += input[i];
    }

    for (int i = 0; i < size; i++) {
        input[i] /= sum;
    }
}

static void forward(NeuralNetwork* nn, float* input, float* output) {
    float hidden1[HIDDEN_SIZE1];
    float hidden2[HIDDEN_SIZE2];

    // Input to first hidden layer
    for (int i = 0; i < HIDDEN_SIZE1; i++) {
        hidden1[i] = 0;
        for (int j = 0; j < INPUT_SIZE; j++) {
            hidden1[i] += input[j] * nn->input_weights[j][i];
        }
        hidden1[i] = relu(hidden1[i] + nn->hidden_bias1[i]);
    }

    // First hidden layer to second hidden layer
    for (int i = 0; i < HIDDEN_SIZE2; i++) {
        hidden2[i] = 0;
        for (int j = 0; j < HIDDEN_SIZE1; j++) {
            hidden2[i] += hidden1[j] * nn->hidden_weights1[j][i];
        }
        hidden2[i] = relu(hidden2[i] + nn->hidden_bias2[i]);
    }

    // Second hidden layer to output
    for (int i = 0; i < OUTPUT_SIZE; i++) {
        output[i] = 0;
        for (int j = 0; j < HIDDEN_SIZE2; j++) {
            output[i] += hidden2[j] * nn->hidden_weights2[j][i];
        }
        output[i] += nn->output_bias[i];
    }

    // Apply softmax to output
    softmax(output, OUTPUT_SIZE);
}

uint32_t predict_color(NeuralNetwork* nn, float red, float green, float blue, float clear) {
    float input[INPUT_SIZE] = {red / 2048.0f, green / 2048.0f, blue / 2048.0f, clear / 2048.0f};
    float output[OUTPUT_SIZE];

    forward(nn, input, output);

             ESP_LOGD(TAG, "Predicted color probabilities:");
            ESP_LOGD(TAG, "Red: %.2f", output[0]);
            ESP_LOGD(TAG, "Black: %.2f", output[1]);
            ESP_LOGD(TAG, "Green: %.2f", output[2]);
            ESP_LOGD(TAG, "White: %.2f", output[3]);

            const char* color_names[] = {"Red", "Black", "Green", "White"};
            uint32_t max_index = 0;
            for (int i = 1; i < OUTPUT_SIZE; i++) {
                if (output[i] > output[max_index]) {
                    max_index = i;
                }
            }
            ESP_LOGW(TAG, "Predicted color: %s", color_names[max_index]);

            return max_index;
            }

            void initialize_neural_network(NeuralNetwork* nn) {
                // Replace these placeholder values with the actual values generated by the Python script
                uint32_t input_weights[INPUT_SIZE][HIDDEN_SIZE1] = {
                        {0x3dce7ff6, 0x401e01c2, 0x3f459587, 0x400acd7e, 0xbf93b293, 0x3f2cf3d9, 0x3dbe0fee, 0xbe966725, 0xbe3eca2c, 0xbfb2791a, 0xbd73893e, 0xc00ff0f9, 0x3e6a0ceb, 0x3f3c1ffe, 0x3f152aea, 0xbe1e4345},
                        {0xbe871da1, 0xbfa83525, 0x3ff85ca0, 0x3ee6bebb, 0x3e3af972, 0xbe95efb6, 0x3f06985c, 0x3fd70757, 0xbf87f079, 0x3f106880, 0x3fa1f1f0, 0x3f40f90e, 0xbfc2e236, 0x3bf3808c, 0xc00b8537, 0xbfc3463b},
                        {0x3e4727eb, 0x3fb70b61, 0x3f187bcc, 0x3f89ce30, 0xbf77adbf, 0x3fd85ae7, 0xbc9ef28e, 0x3f915603, 0xbedd5458, 0xbd7e8b58, 0xbf76fde8, 0xbf8105a0, 0xbfa62a86, 0xbd80ff74, 0xbfbf5b8a, 0xbf8c73c0},
                        {0xbeb4bb83, 0x3e755932, 0xbda6a6a8, 0xbca989d1, 0xbdcd8bf7, 0x3e90b015, 0xbf4f5c9b, 0x3f2422c6, 0x3f038e7b, 0x3faccd13, 0x3e0d272d, 0x3ff347d2, 0xbf0afc03, 0xbf162bd6, 0x3f8f54f9, 0xbee48cae},
                };


                uint32_t hidden_weights1[HIDDEN_SIZE1][HIDDEN_SIZE2] = {
                        {0xbeddbd91, 0x3f6c2506, 0xbd911552, 0xbe76f383, 0x3ed556cd, 0x3da8e90b, 0x3e49ea95, 0xbeb2d561},
                        {0xbf1a278d, 0x3ca261d7, 0x3fedee4f, 0xbfb821c0, 0xbe097b23, 0xbe7e2793, 0x3fab0160, 0xbee746cd},
            {0x3f8c4ec1, 0xbede6a3f, 0x3fc49ba2, 0x3e48b249, 0x3ebb0e0d, 0x3f37e7e8, 0xbf485926, 0x3e7d986a},
            {0xbecc1e5e, 0x3cb84720, 0x3fd03b59, 0xbf795731, 0x3f3fdb72, 0x3edd316c, 0xbe94472d, 0xbe02ab59},
            {0xbe5c034e, 0x3f404120, 0x3f1605c2, 0x3e5eb4f7, 0xbe289aa5, 0x3c2ddfd5, 0xbe87485e, 0xbf22705f},
            {0xbed082a6, 0x3ec899c8, 0x3f5cc247, 0xbdefa97d, 0x3f34ec47, 0xbec36f74, 0xbed7b3cc, 0x3c916d35},
            {0xbdf4806f, 0xbecec7c0, 0xbc73454b, 0x3dfdec3e, 0x3ead244d, 0x3dbaa0f1, 0x3f4ed45e, 0xbdfef530},
            {0x3fd68018, 0xbcb75ceb, 0x3f8dee75, 0xbae6561d, 0xbd807fb9, 0x3df4d7c4, 0xbee11a0a, 0xbdaf9fda},
            {0xbf336719, 0x3c49d37b, 0xbf2c21c1, 0x3e81ad0b, 0x3ddc0c19, 0xbec2aa5b, 0x3f29ba8d, 0xbda87acc},
            {0x3f3f73a8, 0xbf4b89b5, 0xbd185663, 0x3f686f07, 0xbe27f74f, 0xbe28d1fd, 0x3cb4593c, 0x3e810361},
            {0x3f3eef19, 0x3f34aac4, 0x3de40d11, 0xbd210d3d, 0x3c3d2d8a, 0x3e374f9b, 0xbe94b745, 0x3c7a17ab},
            {0x3fc60b9a, 0xbe8e5dd0, 0xbf5e80ec, 0x3f803072, 0xbde40889, 0xbd405f1b, 0x3e839d7f, 0xbe986ca2},
            {0x3e5a3297, 0xbeebd1e2, 0xbd293f51, 0x3e872d1e, 0x3f46442f, 0x3f5b6e08, 0x3f0ba50d, 0x3e823e3c},
            {0x3ee984dd, 0x3db408e0, 0x3e809863, 0xbe9ccc61, 0xbea47da9, 0x3e63df62, 0x3f27a76d, 0xba9d4194},
            {0xbfd16413, 0xbd3cc8e7, 0xbe818b79, 0xbe8174c0, 0xbe35a6aa, 0x3e5ae9db, 0x4006262e, 0xbdb2febc},
            {0x3df66350, 0xbe836eca, 0xbe5d5139, 0xbd7b160e, 0xbbf3ebbc, 0xbf4f0ecc, 0x3e8ccde4, 0x3d26607e},
    };

    uint32_t hidden_weights2[HIDDEN_SIZE2][OUTPUT_SIZE] = {
            {0xbec535c5, 0xc0175ba4, 0x400e5d6e, 0xbf832f19},
            {0x3eda3cf3, 0xbd8a0d44, 0xbf0ae5c1, 0x3ef72346},
            {0x3fa4f9a5, 0xc00bb0ce, 0xbfdd5d0a, 0x3ff993c4},
            {0xbfb11a1a, 0x3e412ab1, 0x3fc9744c, 0xbe8ab1d8},
            {0xbdfd7557, 0xbf07ad4e, 0xbf8cc4c5, 0x3f319475},
            {0x3f255c30, 0xbf4fbc1d, 0x3ed62111, 0x3e0a3a56},
            {0x3f405f36, 0x401f25bd, 0xbf9aabb0, 0xbf78920d},
            {0x3ef9266c, 0xbf11d1bb, 0x3f25a439, 0xbe24bc17},
    };


    uint32_t hidden_bias1[HIDDEN_SIZE1] = {0x0, 0x3d26afd5, 0xbee89e2a, 0xbec19fae, 0x0, 0xbe98e479, 0x0, 0xbe172c28,
                                           0x3ea6bccb, 0x3d61e7b9, 0xbdb4ef70, 0x3ebe1c6d, 0x0, 0xbd7032e0, 0x3f8e17ec, 0x0};

    uint32_t hidden_bias2[HIDDEN_SIZE2] = {0x3e2b2fc8, 0x3c82fd64, 0x3f4c8c0c, 0xbf009713, 0xbd5baff6, 0x3dc2def8, 0x3f7d654f, 0x0};

    uint32_t output_bias[OUTPUT_SIZE] = {0x3f1e3c58, 0xbd665b7d, 0x3def0633, 0xbf2db793};

    // Copy the values to the neural network structure
    for (int i = 0; i < INPUT_SIZE; i++) {
        for (int j = 0; j < HIDDEN_SIZE1; j++) {
            nn->input_weights[i][j] = *(float*)&input_weights[i][j];
        }
    }

    for (int i = 0; i < HIDDEN_SIZE1; i++) {
        for (int j = 0; j < HIDDEN_SIZE2; j++) {
            nn->hidden_weights1[i][j] = *(float*)&hidden_weights1[i][j];
        }
    }

    for (int i = 0; i < HIDDEN_SIZE2; i++) {
        for (int j = 0; j < OUTPUT_SIZE; j++) {
            nn->hidden_weights2[i][j] = *(float*)&hidden_weights2[i][j];
        }
    }

    for (int i = 0; i < HIDDEN_SIZE1; i++) {
        nn->hidden_bias1[i] = *(float*)&hidden_bias1[i];
    }

    for (int i = 0; i < HIDDEN_SIZE2; i++) {
        nn->hidden_bias2[i] = *(float*)&hidden_bias2[i];
    }

    for (int i = 0; i < OUTPUT_SIZE; i++) {
        nn->output_bias[i] = *(float*)&output_bias[i];
    }
}

